import {
  addEdge,
  applyEdgeChanges,
  applyNodeChanges,
  Background,
  Controls,
  Handle,
  ReactFlow,
  Position,
  getBezierPath,
  BaseEdge
} from "@xyflow/react";
import type { Edge, EdgeChange, Node, NodeChange, Connection, NodeProps, NodeTypes, EdgeProps } from "@xyflow/react";
import "@xyflow/react/dist/style.css";
import { memo, useCallback, useEffect, useRef, useState } from "react";

import * as d3 from "d3-force";

export type CustomNode = Node<{ label: string }>;

// 自定义节点
const CustomNode = memo((props: NodeProps<CustomNode>) => {
  return (
    <div className="bg-amber-200 p-4 rounded shadow">
      <Handle type="source" position={Position.Bottom} />
      <div className="text-lg font-bold">{props.data.label}</div>
      <Handle type="target" position={Position.Top} />
    </div>
  );
});

const nodeTypes: NodeTypes = {
  customNode: CustomNode
};

// 自定义边连线
const CustomEdge = ({ sourceX, sourceY, targetX, targetY }: EdgeProps) => {
  // 连线的方式
  // 主要有三种连线的方式：
  // 直线(getStraightPath)、曲线(getSmoothStepPath)、贝塞尔曲线(getBezierPath)
  const [edgePath] = getBezierPath({
    sourceX,
    sourceY,
    targetX,
    targetY
  });

  // 使用默认的 BaseEdge 组件来渲染边
  // return <BaseEdge path={edgePath} />;

  // 如果需要自定义边的样式，可以直接返回 SVG 路径
  return (
    <>
      <path d={edgePath} stroke="red" strokeWidth={3} fill="none" strokeDasharray="5,5" />
    </>
  );
};

const edgeTypes = {
  customEdge: CustomEdge
};

const initialNodes = [
  {
    id: "n1",
    type: "customNode",
    data: { label: "Input Node" },
    position: { x: 100, y: 50 }
  },
  {
    id: "n2",
    type: "customNode",
    data: { label: "Output Node1" },
    position: { x: 100, y: 200 }
  },
  {
    id: "n3",
    type: "customNode",
    data: { label: "Output Node2" },
    position: { x: 300, y: 200 }
  },
  {
    id: "n4",
    type: "customNode",
    data: { label: "Output Node3" },
    position: { x: 500, y: 200 }
  }
];

const initialEdges: Edge[] = [
  // 都连接到n1节点
  // { id: "n1-n2", source: "n1", target: "n2", type: "customEdge" },
  // { id: "n1-n3", source: "n1", target: "n3", type: "customEdge" },
  // { id: "n1-n4", source: "n1", target: "n4", type: "customEdge" }
  // 都连接到n2节点
  { id: "n2-n1", source: "n2", target: "n1", type: "customEdge" },
  { id: "n2-n3", source: "n2", target: "n3", type: "customEdge" },
  { id: "n2-n4", source: "n2", target: "n4", type: "customEdge" }
];

export function DocGraph() {
  const [nodes, setNodes] = useState<Node[]>(initialNodes);
  const [edges, setEdges] = useState<Edge[]>(initialEdges);

  const simulationRef = useRef<d3.Simulation<any, any> | null>(null); // 🔧 保存 d3 simulation 实例
  const nodeMapRef = useRef<Map<string, any>>(new Map());

  const draggingNodeIdRef = useRef<string | null>(null);

  const onNodesChange = useCallback(
    (changes: NodeChange[]) => setNodes(nodesSnapshot => applyNodeChanges(changes, nodesSnapshot)),
    []
  );

  const onEdgesChange = useCallback(
    (changes: EdgeChange[]) => setEdges(edgesSnapshot => applyEdgeChanges(changes, edgesSnapshot)),
    []
  );

  const onConnect = useCallback((params: Connection) => setEdges(edgesSnapshot => addEdge(params, edgesSnapshot)), []);

  // ✅ 监听拖拽时的节点位置并更新到 d3 simulation
  const onNodeDrag = useCallback((event, node) => {
    draggingNodeIdRef.current = node.id;
    const simNode = nodeMapRef.current.get(node.id);
    if (simNode) {
      simNode.fx = node.position.x;
      simNode.fy = node.position.y;
      simulationRef.current?.alpha(0.3).restart();
    }
  }, []);

  // ✅ 拖拽完成后也保留节点位置（你可以改为释放力场）
  const onNodeDragStop = useCallback((event, node) => {
    draggingNodeIdRef.current = null;
    const simNode = nodeMapRef.current.get(node.id);
    if (simNode) {
      simNode.fx = node.position.x;
      simNode.fy = node.position.y;
    }
  }, []);

  useEffect(() => {
    const nodeMap = new Map(nodes.map(n => [n.id, n]));
    nodeMapRef.current = nodeMap;

    const d3Links = edges.map(edge => ({
      source: nodeMap.get(edge.source), // 从 ID 映射为 node 对象
      target: nodeMap.get(edge.target)
    }));
    // 使用 D3 力导向布局来自动调整节点位置
    const simulation = d3
      .forceSimulation(nodes)
      .force("charge", d3.forceManyBody().strength(-200)) // 更强的排斥力，防止挤在一起
      .force("collide", d3.forceCollide().radius(200)) // 冲突检测，保持间距
      .force(
        "link",
        d3
          .forceLink(d3Links)
          .id((d: any) => d.id)
          .distance(150)
          .strength(0.8)
      ) // 线更紧凑
      .force("center", d3.forceCenter(700, 350)) // 吸向中心
      .force("radial", d3.forceRadial(0, 700, 350)) // 径向力
      .alpha(1) // 初始能量
      .alphaDecay(0.05); // 能量衰减速度，0.05 会让图在 60~80 帧内收敛

    simulationRef.current = simulation;

    simulation.on("end", () => {
      console.log("end");
    });

    simulation.on("tick", () => {
      const forceNodes = simulation.nodes();
      const draggingId = draggingNodeIdRef.current;

      const newNodes = forceNodes
        .map((node: any) => {
          // 跳过被拖动的节点，保留其当前 ReactFlow 位置不变
          if (node.id === draggingId) return null;

          return {
            ...node,
            position: { x: node.x, y: node.y }
          };
        })
        .filter(Boolean); // 去除 null

      setNodes(prevNodes =>
        prevNodes.map(n => {
          const updated = newNodes.find(nn => nn.id === n.id);
          return updated ? { ...n, position: updated.position } : n;
        })
      );

      // const forceNodes = simulation.nodes();
      // const newNodes = forceNodes.map((node: any) => ({
      //   ...node,
      //   position: { x: node.x, y: node.y }
      // }));
      // setNodes(newNodes);
    });

    return () => {
      simulation.stop();
    };
  }, [nodes]);

  return (
    <div className="h-screen w-full">
      <ReactFlow
        nodeTypes={nodeTypes}
        edgeTypes={edgeTypes}
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        onNodeDrag={onNodeDrag} // ✅ 拖拽中更新力模型
        onNodeDragStop={onNodeDragStop} // ✅ 拖拽完成后可固定位置
      >
        <Background />
        <Controls />
      </ReactFlow>
    </div>
  );
}
